import numpy as np
import heapq
import random
import math
from collections import deque

class RRTPathPlanner:
    """RRT (Rapidly-exploring Random Tree) 路径规划算法"""
    
    def __init__(self, grid_size, max_iterations=1000, step_size=20, goal_bias=0.1):
        self.grid_size = grid_size
        self.max_iterations = max_iterations
        self.step_size = step_size
        self.goal_bias = goal_bias
        self.safe_distance = 25
        
    def distance(self, p1, p2):
        """计算两点间距离"""
        return np.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)
    
    def is_collision_free(self, p1, p2, obstacles):
        """检查路径段是否无碰撞"""
        # 沿路径检查多个点
        num_checks = int(self.distance(p1, p2) / 5) + 1
        for i in range(num_checks + 1):
            t = i / num_checks if num_checks > 0 else 0
            check_point = (
                p1[0] + t * (p2[0] - p1[0]),
                p1[1] + t * (p2[1] - p1[1])
            )
            
            # 边界检查
            if (check_point[0] < 0 or check_point[0] >= self.grid_size or 
                check_point[1] < 0 or check_point[1] >= self.grid_size):
                return False
            
            # 障碍物检查
            for obs in obstacles:
                if self.distance(check_point, obs) < self.safe_distance:
                    return False
        return True
    
    def get_random_point(self):
        """获取随机采样点"""
        return (random.uniform(0, self.grid_size), random.uniform(0, self.grid_size))
    
    def find_nearest_node(self, tree, point):
        """找到树中最近的节点"""
        min_dist = float('inf')
        nearest_node = None
        for node in tree:
            dist = self.distance(node['point'], point)
            if dist < min_dist:
                min_dist = dist
                nearest_node = node
        return nearest_node
    
    def steer(self, from_point, to_point):
        """从from_point向to_point方向扩展step_size距离"""
        dist = self.distance(from_point, to_point)
        if dist <= self.step_size:
            return to_point
        
        theta = math.atan2(to_point[1] - from_point[1], to_point[0] - from_point[0])
        new_point = (
            from_point[0] + self.step_size * math.cos(theta),
            from_point[1] + self.step_size * math.sin(theta)
        )
        return new_point
    
    def find_path(self, start, goal, obstacles):
        """RRT路径搜索"""
        start = (float(start[0]), float(start[1]))
        goal = (float(goal[0]), float(goal[1]))
        
        # 初始化树
        tree = [{'point': start, 'parent': None}]
        
        for _ in range(self.max_iterations):
            # 采样点（带目标偏置）
            if random.random() < self.goal_bias:
                sample_point = goal
            else:
                sample_point = self.get_random_point()
            
            # 找最近节点
            nearest_node = self.find_nearest_node(tree, sample_point)
            
            # 扩展
            new_point = self.steer(nearest_node['point'], sample_point)
            
            # 碰撞检测
            if self.is_collision_free(nearest_node['point'], new_point, obstacles):
                new_node = {'point': new_point, 'parent': nearest_node}
                tree.append(new_node)
                
                # 检查是否到达目标
                if self.distance(new_point, goal) < self.step_size:
                    # 构建路径
                    path = []
                    current = new_node
                    while current is not None:
                        path.append(current['point'])
                        current = current['parent']
                    path.reverse()
                    return path
        
        # 如果找不到路径，返回直线路径
        return [start, goal]


class PotentialFieldPlanner:
    """人工势场法路径规划"""
    
    def __init__(self, grid_size, attractive_gain=1.0, repulsive_gain=100.0, 
                 repulsive_range=50.0, step_size=10.0):
        self.grid_size = grid_size
        self.attractive_gain = attractive_gain
        self.repulsive_gain = repulsive_gain
        self.repulsive_range = repulsive_range
        self.step_size = step_size
        self.max_iterations = 200
        
    def attractive_force(self, current, goal):
        """计算引力"""
        dx = goal[0] - current[0]
        dy = goal[1] - current[1]
        distance = np.sqrt(dx**2 + dy**2)
        
        if distance < 1e-6:
            return np.array([0.0, 0.0])
        
        # 线性引力
        force_magnitude = self.attractive_gain * distance
        force_x = force_magnitude * dx / distance
        force_y = force_magnitude * dy / distance
        
        return np.array([force_x, force_y])
    
    def repulsive_force(self, current, obstacles):
        """计算斥力"""
        total_force = np.array([0.0, 0.0])
        
        for obstacle in obstacles:
            dx = current[0] - obstacle[0]
            dy = current[1] - obstacle[1]
            distance = np.sqrt(dx**2 + dy**2)
            
            if distance < self.repulsive_range and distance > 1e-6:
                # 反比例斥力
                force_magnitude = self.repulsive_gain * (1.0/distance - 1.0/self.repulsive_range) / (distance**2)
                force_x = force_magnitude * dx / distance
                force_y = force_magnitude * dy / distance
                total_force += np.array([force_x, force_y])
        
        return total_force
    
    def boundary_force(self, current):
        """边界斥力"""
        force = np.array([0.0, 0.0])
        boundary_margin = 30
        
        # 左边界
        if current[0] < boundary_margin:
            force[0] += self.repulsive_gain / (current[0] + 1)
        # 右边界
        if current[0] > self.grid_size - boundary_margin:
            force[0] -= self.repulsive_gain / (self.grid_size - current[0] + 1)
        # 下边界
        if current[1] < boundary_margin:
            force[1] += self.repulsive_gain / (current[1] + 1)
        # 上边界
        if current[1] > self.grid_size - boundary_margin:
            force[1] -= self.repulsive_gain / (self.grid_size - current[1] + 1)
            
        return force
    
    def find_path(self, start, goal, obstacles):
        """势场法路径规划"""
        path = [start]
        current = np.array(start, dtype=float)
        goal_array = np.array(goal, dtype=float)
        
        for _ in range(self.max_iterations):
            # 计算总力
            attractive = self.attractive_force(current, goal_array)
            repulsive = self.repulsive_force(current, obstacles)
            boundary = self.boundary_force(current)
            
            total_force = attractive + repulsive + boundary
            
            # 归一化并移动
            force_magnitude = np.linalg.norm(total_force)
            if force_magnitude > 1e-6:
                direction = total_force / force_magnitude
                next_pos = current + direction * self.step_size
                
                # 边界限制
                next_pos[0] = np.clip(next_pos[0], 0, self.grid_size)
                next_pos[1] = np.clip(next_pos[1], 0, self.grid_size)
                
                current = next_pos
                path.append(tuple(current))
                
                # 检查是否到达目标
                if np.linalg.norm(current - goal_array) < 15:
                    break
            else:
                break
        
        return path


class DijkstraPlanner:
    """Dijkstra最短路径算法"""
    
    def __init__(self, grid_size, resolution=10):
        self.grid_size = grid_size
        self.resolution = resolution
        self.grid_width = grid_size // resolution
        self.grid_height = grid_size // resolution
        
    def world_to_grid(self, world_pos):
        """世界坐标转网格坐标"""
        return (int(world_pos[0] // self.resolution), int(world_pos[1] // self.resolution))
    
    def grid_to_world(self, grid_pos):
        """网格坐标转世界坐标"""
        return (grid_pos[0] * self.resolution + self.resolution//2, 
                grid_pos[1] * self.resolution + self.resolution//2)
    
    def is_valid_grid(self, grid_pos):
        """检查网格位置是否有效"""
        return (0 <= grid_pos[0] < self.grid_width and 
                0 <= grid_pos[1] < self.grid_height)
    
    def create_obstacle_map(self, obstacles):
        """创建障碍物地图"""
        obstacle_map = np.zeros((self.grid_width, self.grid_height), dtype=bool)
        
        for obs in obstacles:
            grid_obs = self.world_to_grid(obs)
            # 在障碍物周围设置安全区域
            for dx in range(-2, 3):
                for dy in range(-2, 3):
                    new_pos = (grid_obs[0] + dx, grid_obs[1] + dy)
                    if self.is_valid_grid(new_pos):
                        obstacle_map[new_pos[0], new_pos[1]] = True
        
        return obstacle_map
    
    def get_neighbors(self, pos):
        """获取8邻域"""
        neighbors = []
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                if dx == 0 and dy == 0:
                    continue
                new_pos = (pos[0] + dx, pos[1] + dy)
                if self.is_valid_grid(new_pos):
                    # 对角线移动成本更高
                    cost = 1.414 if abs(dx) + abs(dy) == 2 else 1.0
                    neighbors.append((new_pos, cost))
        return neighbors
    
    def find_path(self, start, goal, obstacles):
        """Dijkstra路径搜索"""
        start_grid = self.world_to_grid(start)
        goal_grid = self.world_to_grid(goal)
        
        obstacle_map = self.create_obstacle_map(obstacles)
        
        # 如果起点或终点在障碍物上，返回直线路径
        if (obstacle_map[start_grid[0], start_grid[1]] or 
            obstacle_map[goal_grid[0], goal_grid[1]]):
            return [start, goal]
        
        # Dijkstra算法
        distances = {start_grid: 0}
        previous = {}
        pq = [(0, start_grid)]
        visited = set()
        
        while pq:
            current_dist, current = heapq.heappop(pq)
            
            if current in visited:
                continue
            
            visited.add(current)
            
            if current == goal_grid:
                # 重构路径
                path = []
                while current in previous:
                    path.append(self.grid_to_world(current))
                    current = previous[current]
                path.append(self.grid_to_world(start_grid))
                path.reverse()
                return path
            
            for neighbor, move_cost in self.get_neighbors(current):
                if neighbor in visited or obstacle_map[neighbor[0], neighbor[1]]:
                    continue
                
                new_dist = current_dist + move_cost
                
                if neighbor not in distances or new_dist < distances[neighbor]:
                    distances[neighbor] = new_dist
                    previous[neighbor] = current
                    heapq.heappush(pq, (new_dist, neighbor))
        
        # 找不到路径，返回直线
        return [start, goal]


class HybridPathPlanner:
    """混合路径规划器 - 根据情况选择最适合的算法"""
    
    def __init__(self, grid_size):
        self.grid_size = grid_size
        self.astar = None  # 将在需要时初始化
        self.rrt = RRTPathPlanner(grid_size)
        self.potential_field = PotentialFieldPlanner(grid_size)
        self.dijkstra = DijkstraPlanner(grid_size)
        
    def analyze_environment(self, start, goal, obstacles):
        """分析环境特征，选择最适合的算法"""
        distance_to_goal = np.linalg.norm(np.array(goal) - np.array(start))
        num_obstacles = len(obstacles)
        
        # 计算障碍物密度
        obstacle_density = num_obstacles / (self.grid_size ** 2) * 10000
        
        # 检查是否有障碍物在起点和终点之间
        obstacles_in_path = 0
        for obs in obstacles:
            # 简单的线段-点距离检查
            if self._point_to_line_distance(obs, start, goal) < 30:
                obstacles_in_path += 1
        
        return {
            'distance': distance_to_goal,
            'obstacle_count': num_obstacles,
            'obstacle_density': obstacle_density,
            'obstacles_in_path': obstacles_in_path
        }
    
    def _point_to_line_distance(self, point, line_start, line_end):
        """计算点到线段的距离"""
        A = np.array(line_start)
        B = np.array(line_end)
        P = np.array(point)
        
        AB = B - A
        AP = P - A
        
        if np.dot(AB, AB) == 0:
            return np.linalg.norm(AP)
        
        t = np.dot(AP, AB) / np.dot(AB, AB)
        t = max(0, min(1, t))
        
        closest = A + t * AB
        return np.linalg.norm(P - closest)
    
    def select_algorithm(self, env_analysis):
        """根据环境分析选择算法"""
        distance = env_analysis['distance']
        obstacle_density = env_analysis['obstacle_density']
        obstacles_in_path = env_analysis['obstacles_in_path']
        
        # 决策逻辑
        if obstacle_density < 0.1 and obstacles_in_path == 0:
            return 'potential_field'  # 简单环境用势场法
        elif distance < 100 and obstacle_density < 0.3:
            return 'astar'  # 中等复杂度用A*
        elif obstacle_density > 0.5 or obstacles_in_path > 3:
            return 'rrt'  # 复杂环境用RRT
        else:
            return 'dijkstra'  # 其他情况用Dijkstra
    
    def find_path(self, start, goal, obstacles):
        """混合路径规划"""
        env_analysis = self.analyze_environment(start, goal, obstacles)
        algorithm = self.select_algorithm(env_analysis)
        
        try:
            if algorithm == 'potential_field':
                return self.potential_field.find_path(start, goal, obstacles)
            elif algorithm == 'rrt':
                return self.rrt.find_path(start, goal, obstacles)
            elif algorithm == 'dijkstra':
                return self.dijkstra.find_path(start, goal, obstacles)
            else:  # astar
                # 如果需要A*，使用现有的实现
                if self.astar is None:
                    from hierarchical_agent import AStarPathPlanner
                    self.astar = AStarPathPlanner(self.grid_size)
                return self.astar.find_path(start, goal, obstacles)
        except Exception as e:
            print(f"路径规划算法 {algorithm} 失败: {e}")
            # 降级到简单的直线路径
            return [start, goal]